# import torch
import joblib
from sklearn import tree


class DecisionTree:
    '''
    决策树,输入用户金融数据，输出用户的分类型, 使用CART决策算法
    '''
    data_set = {}  # 数据集
    example_cnt_threshold = 0  # 样本数量阈值
    gini_threshold = 0  # 基尼系数阈值

    def __init__(self, data_set, example_cnt_threshold, gini_threshold):
        self.data_set = data_set
        self.example_cnt_threshold = example_cnt_threshold
        self.gini_threshold = gini_threshold

    def create_tree(self, data_set, lables):

        '''
        :param data_set:需要构建决策树的数据集
        :param self: 决策树
        :return:
        '''
        self.data_set = data_set
        # 获取数据集的标签列表（最后一列）
        class_list = [example[-1] for example in data_set]
        # 终止条件1：如果数据集的所有标签都相同，则返回该标签
        if class_list.count(class_list[0]) == len(class_list):
            return class_list[0]
        # 终止条件2：todo

        pass

    def classify(self, input_tree, feat_labels, test_vec):
        '''
        :param input_tree: 决策树
        :param feat_labels: 特征标签
        :param test_vec: 测试数据
        :return:
        '''
        first_str = list(input_tree.keys())[0]
        second_dict = input_tree[first_str]
        feat_index = feat_labels.index(first_str)
        key = test_vec[feat_index]
        value_of_feat = second_dict[key]
        if isinstance(value_of_feat, dict):
            class_label = self.classify(value_of_feat, feat_labels, test_vec)
        else:
            class_label = value_of_feat
        return class_label

    # sklearn 的决策树


def classify():
    '''
    test_classifier
    :return: predict result and predict_proba
    '''
    # X: train[N_SAMPLES, N_FEATURES]
    # Y: label[N_SAMPLES]
    X = [[0, 0], [1, 1], [0, 1], [2, 2], [3, 3], [1, 2], [2, 1], [1, 0], [0, 1]]
    Y = [0, 1, 1, 2, 7, 4, 1, 2, 0]
    clf = tree.DecisionTreeClassifier(criterion='gini', splitter='best', max_depth=10)
    clf = clf.fit(X, Y)

    # predict
    return clf.predict([[3, 3]]), clf.predict_proba([[3, 3]])



def classify_with_multi_features():
    '''
    test_classifier：大于三特征值的分类
    :return: predict result and predict_proba
    '''
    X_train= [[0, 0, 0], [1, 1, 1], [0, 1, 1], [2, 2, 2], [3, 3, 3], [1, 2, 2], [2, 1, 1], [1, 0, 0], [0, 1, 1]]
    Y_train = [0, 1, 1, 2, 7, 4, 1, 2, 0]
    clf = tree.DecisionTreeClassifier(criterion='gini', splitter='best', max_depth=10)
    clf.fit(X_train, Y_train)

    joblib.dump(clf, 'tree.pkl')
    return clf.predict([[1, 1, 2]]), clf.predict_proba([[1, 1, 2]])


# get an idea
# 是否可以将持仓类型作为tag，以proba的某种变异形式作为推荐持仓结果

def classify_with_trained_modle():
    # load modle
    clf = joblib.load('tree.pkl')
    return clf.predict([[1, 1, 2]]), clf.predict_proba([[1, 1, 2]])



if __name__ == '__main__':
    print(classify_with_multi_features())
    print(classify())